import os,sys
import tensorflow as tf
from PIL import Image
import matplotlib.pyplot as plt


def tf_visualization(data_path):
    swd = os.path.basename(data_path)
    if os.path.exists(swd) is False:
        os.makedirs(swd)

    data_files = tf.gfile.Glob(data_path)
    if len(data_files) == 0:
        print(" no tf file avaliable")
        exit(-1)

    filename_queue = tf.train.string_input_producer(data_files, shuffle=True)
    reader = tf.TFRecordReader()
    _, serialized_example = reader.read(filename_queue)
    features = tf.parse_single_example(serialized_example,
                                       features={
                                           'label': tf.FixedLenFeature([], tf.string),
                                           'img_raw': tf.FixedLenFeature([], tf.string),
                                           'img_width': tf.FixedLenFeature([], tf.int64),
                                           'img_height': tf.FixedLenFeature([], tf.int64),
                                       })
    image = tf.decode_raw(features['img_raw'], tf.uint8)
    height = tf.cast(features['img_height'], tf.int32)
    width = tf.cast(features['img_width'], tf.int32)
    label = tf.cast(features['label'], tf.string)

    # label = bytes(label)
    # label = bytes(label).decode(encoding="utf-8")
    channel = 3
    image = tf.reshape(image, [height, width, channel])

    head_img = {}
    head_img_inited = False
    i = 1

    with tf.Session() as sess:
        init_op = tf.initialize_all_variables()
        sess.run(init_op)
        coord = tf.train.Coordinator()
        threads = tf.train.start_queue_runners(coord=coord)
        while True:
            single, l = sess.run([image, label])
            l = l.decode()
            img = Image.fromarray(single, 'RGB')
            if img == head_img:
                break
            if head_img_inited is False:
                head_img = img
                head_img_inited = True
            save_path = swd + "/" + str(l) + "_" + str(i) + '.jpg'
            print("saving to ", save_path)
            img.save(save_path)
            i = i + 1
            # print(single,l)
        coord.request_stop()
        coord.join(threads)

if __name__ == "__main__":
    if len(sys.argv) < 3:
        print("usage: ./classifier_tf_vis.py -i xxx.tfrecords")
        exit(-1)
    data_path = sys.argv[2]
    tf_visualization(data_path)
